package com.benromberg.junit.threadconductor;

import java.lang.reflect.Method;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicInteger;

import org.junit.runner.Description;
import org.junit.runner.notification.RunNotifier;
import org.junit.runners.BlockJUnit4ClassRunner;
import org.junit.runners.model.FrameworkMethod;
import org.junit.runners.model.InitializationError;
import org.junit.runners.model.Statement;

public class ConcurrentTestRunner extends BlockJUnit4ClassRunner {
	private Thread mainThread = Thread.currentThread();
	private int numThreads = 10;
	private int timeout = 0;
	private AtomicInteger waitCount;
	private Throwable failure;

	private boolean runIgnored = false;

	@Override
	protected void runChild(final FrameworkMethod method, RunNotifier notifier) {
		if (!runIgnored) {
			super.runChild(method, notifier);
		} else {
			Description description = describeChild(method);
			runLeaf(methodBlock(method), description, notifier);
		}
	}

	public void setRunIgnored(boolean runIgnored) {
		this.runIgnored = runIgnored;
	}

	public ConcurrentTestRunner(Class<?> klass) throws InitializationError {
		super(klass);
	}

	private synchronized void setFailure(Throwable failure) {
		this.failure = failure;
	}

	private Statement superMethodBlock(final FrameworkMethod method) {
		return super.methodBlock(method);
	}

	private void initFromAnnotations(final FrameworkMethod method) {
		Method testMethod = method.getMethod();
		ConcurrentTest classAnnotation = testMethod.getDeclaringClass()
				.getAnnotation(ConcurrentTest.class);
		if (classAnnotation != null) {
			numThreads = classAnnotation.numThreads();
			timeout = classAnnotation.timeout();
		}
		ConcurrentTestMethod methodAnnotation = testMethod
				.getAnnotation(ConcurrentTestMethod.class);
		if (methodAnnotation != null) {
			numThreads = methodAnnotation.numThreads();
			timeout = methodAnnotation.timeout();
		}
		waitCount = new AtomicInteger(numThreads);
	}

	private void startThreads(final FrameworkMethod method) {
		System.out.println("starting threads");
		for (int i = 0; i < numThreads; i++) {
			new Thread(new Runnable() {
				@Override
				public void run() {
					try {
						superMethodBlock(method).evaluate();
					} catch (Throwable e) {
						setFailure(e);
					} finally {
						if (waitCount.decrementAndGet() == 0) {
							mainThread.interrupt();
						}
					}
				}
			}).start();
		}
	}

	private synchronized void waitAndEvaluateResult() throws TimeoutException,
			Throwable {
		try {
			if (timeout == 0) {
				wait();
			} else {
				Thread.sleep(timeout);
			}
			throw new TimeoutException(
					"Test timed out while waiting for an expected result");
		} catch (InterruptedException expected) {
		}
		if (failure != null) {
			throw failure;
		}
	}

	@Override
	protected Statement methodBlock(final FrameworkMethod method) {
		return new Statement() {
			@Override
			public void evaluate() throws Throwable {
				initFromAnnotations(method);
				startThreads(method);
				waitAndEvaluateResult();
			}
		};
	}

}
